import os
import pydicom
from PIL import Image
import numpy as np
import shutil

def is_dicom_file(file_path):
    """检查文件是否为 DICOM 文件"""
    try:
        with open(file_path, 'rb') as f:
            # 读取前 128 字节（DICOM 前置数据）
            preamble = f.read(128)
            # 读取接下来的 4 个字节（应该是 "DICM"）
            dicom_identifier = f.read(4)
            if dicom_identifier == b'DICM':
                return True
    except Exception as e:
        print(f"无法读取文件 {file_path}: {e}")
    return False

def convert_dicom_to_jpg(dicom_path, output_path):
    """将 DICOM 文件转换为 JPG 图像"""
    try:
        ds = pydicom.dcmread(dicom_path)
        pixel_array = ds.pixel_array
        # 将像素值缩放到 0-255 的范围
        img = (pixel_array - pixel_array.min()) * (255.0 / (pixel_array.max() - pixel_array.min()))
        img = Image.fromarray(img.astype(np.uint8))
        img.save(output_path, format='JPEG')
        print(f"已成功将 {dicom_path} 转换为 {output_path}")
    except Exception as e:
        print(f"转换失败: {e}")

def ensure_directory_exists(directory):
    """确保目录存在，如果不存在则创建"""
    if not os.path.exists(directory):
        os.makedirs(directory)

def process_folder(input_folder, output_folder, output_format='jpg'):
    """遍历文件夹及其子文件夹，处理无扩展名的 DICOM 文件，并将结果保存到新文件夹"""
    for root, dirs, files in os.walk(input_folder):
        for file in files:
            file_path = os.path.join(root, file)
            # 检查文件是否有扩展名
            if '.' not in file:
                if is_dicom_file(file_path):
                    # 计算相对于输入文件夹的相对路径
                    relative_path = os.path.relpath(root, input_folder)
                    # 创建输出文件夹中的相应子文件夹
                    output_subfolder = os.path.join(output_folder, relative_path)
                    ensure_directory_exists(output_subfolder)
                    
                    # 生成新的文件名
                    new_filename = os.path.splitext(file)[0]
                    if output_format.lower() == 'dcm':
                        new_file_path = os.path.join(output_subfolder, f"{new_filename}.dcm")
                        # 复制文件并重命名为 .dcm
                        shutil.copy2(file_path, new_file_path)
                        print(f"已将 {file_path} 复制并重命名为 {new_file_path}")
                    elif output_format.lower() == 'jpg':
                        new_file_path = os.path.join(output_subfolder, f"{new_filename}.jpg")
                        convert_dicom_to_jpg(file_path, new_file_path)

if __name__ == "__main__":
    # 指定要处理的文件夹路径
    input_folder = r'C:/Users/PC/Desktop/IMAGES'
    
    # 指定输出文件夹路径
    output_folder = r'C:/Users/PC/Desktop/output'
    
    # 选择输出格式：'jpg' 或 'dcm'
    output_format = 'jpg'  # 或 'dcm'
    
    # 确保输出文件夹存在
    ensure_directory_exists(output_folder)
    
    # 处理文件夹中的所有无扩展名 DICOM 文件
    process_folder(input_folder, output_folder, output_format)